from abc import ABC, abstractmethod
from typing import Any, Dict, Optional
import logging
from .exceptions import ModelLoadError, InvalidInputError

class BaseModel(ABC):
    """Base class for all AI models in PolyModel Agent"""
    
    def __init__(self, model_name: str, config: Optional[Dict] = None):
        """
        Initialize base model.
        
        Args:
            model_name (str): Name of the model to load
            config (Optional[Dict]): Model configuration parameters
        """
        self.model_name = model_name
        self.config = config or {}
        self.model = None
        self.logger = logging.getLogger(self.__class__.__name__)
    
    @abstractmethod
    async def load(self) -> None:
        """
        Load model into memory.
        
        Raises:
            ModelLoadError: If model loading fails
        """
        pass
    
    @abstractmethod
    async def predict(self, input_data: Any) -> Any:
        """
        Make predictions with the model.
        
        Args:
            input_data: Input data for prediction
            
        Returns:
            Model predictions
            
        Raises:
            InvalidInputError: If input validation fails
            ModelLoadError: If model is not loaded
        """
        if not self.model:
            raise ModelLoadError("Model not loaded")
        
        if not await self.validate_input(input_data):
            raise InvalidInputError(f"Invalid input: {input_data}")
    
    @abstractmethod
    async def validate_input(self, input_data: Any) -> bool:
        """
        Validate input data before prediction.
        
        Args:
            input_data: Input data to validate
            
        Returns:
            bool: True if input is valid, False otherwise
        """
        pass
    
    async def preprocess(self, input_data: Any) -> Any:
        """
        Preprocess input data before prediction.
        
        Args:
            input_data: Raw input data
            
        Returns:
            Preprocessed input data
        """
        return input_data
    
    async def postprocess(self, output_data: Any) -> Any:
        """
        Postprocess model output.
        
        Args:
            output_data: Raw model output
            
        Returns:
            Processed output data
        """
        return output_data
    
    def __repr__(self) -> str:
        return f"{self.__class__.__name__}(model_name='{self.model_name}')"